#include "PowerSpectrum.h"
#include "Int.h"

#include <fstream>
using std::ifstream;
using std::ofstream;
#include <iostream>
#include <sstream>
using std::istringstream;
#include <string>
using std::string;
#include <map>
using std::map;
#include <vector>
using std::vector;

#include "TMV.h"
using tmv::Vector;
using tmv::Matrix;
#include "dbg.h"
using std::endl;

#define DATECODE "0913"

#define NAME "lambda"
// Add NAME to the output names (lnlike.NAME.dat and predlike.NAME.debug)

#define WMAP3

//#define DILLUT  1.05
// Assume dillution error, so multiply each gamma by DILLUT
//#define BMODE +1.0
// +1.0 for add B mode, -1.0 for subract

//For PD or CFRS, change predfile below

//#define STARTAT 1530
//#define STOPAT 23860
//#define Wa
//#define NonFlat
//#define TAKEEVERYN 100
#define LAMBDA

std::ostream* dbgout=0;
bool XDEBUG=false;

#ifdef Wa

#ifdef NonFlat
const char* cmbfile = "NonFlatWWaPrior.dat";
#ifdef TAKEEVERYN
const char* predfile = "pred.nfwwa_100.out";
#else
const char* predfile = "pred.nfwwa.out";
#endif

#else // !NonFlat
const char* cmbfile = "WWaPrior3.dat";
#ifdef TAKEEVERYN
const char* predfile = "pred.wwa3_100.out";
#else
const char* predfile = "pred.wwa3.out";
#endif
#endif // NonFlat

#else // !Wa
#ifdef LAMBDA
//const char* cmbfile = "LCDMWMAP.dat";
const char* cmbfile = "wmap3_lcdm.dat";
const char* predfile = "pred.lambda.out";
#else
const char* cmbfile = "wmap3_wcdm_pert.dat";
//const char* cmbfile = "CMBQuintchainNew.dat";
//const char* cmbfile = "CMBQuintchain.dat";
const char* predfile = "pred.wmap3.out";
//const char* predfile = "pred.cmb.out";
//const char* predfile = "pred.pd.out";
//const char* predfile = "pred.cfrs.out";
#endif //LAMBDA
#endif //Wa

const double PI = 3.141592653589793;
const double DEG = PI/180.;
const double ARCSEC = DEG/3600.;
const double ARCMIN = DEG/60.;

const double map_minr = 1.0*ARCMIN;
//const double map_minr = 5.0*ARCMIN;
const double map_maxr = 40.0*ARCMIN;
const double var_minr = 40.0*ARCMIN;
const double var_maxr = 100.0*ARCMIN;
const double singthresh = 0.01;
const int datafrac = 1; // take every nth line of the data

#ifdef DILLUT
const double DILLUTSQ = DILLUT*DILLUT;
#endif

void ReadData(Vector<double>*& r, Vector<double>*& mean, Matrix<double>*& cov, 
    size_t& nmapout)
{
  dbg<<"Start ReadData\n";

  const size_t nfields = 12;
  const string datafile = "/home/mjarvis/plots/m2.out.";
  const char* fieldcodes[nfields] = { 
    "A", "G", "H", "J", "K", "L", "M", "N", "Q", "R", "T", "X" };

  vector<vector<double> > rmap(nfields);
  vector<vector<double> > map(nfields);
  //vector<vector<double> > sigmap(nfields);
  vector<vector<double> > rvar(nfields);
  vector<vector<double> > var(nfields);
  //vector<vector<double> > sigvar(nfields);

  for(size_t k=0;k<nfields;k++) {
    string name = datafile + DATECODE + '.' + fieldcodes[k];
    ifstream fin(name.c_str());
    if (!fin) myerror("opening data file:  ",name.c_str());

    dbg<<"Reading file: "<<name<<endl;
    string line;
    int count=0;
    while (getline(fin,line)) {
      if (line[0] == '#') continue;
      count++;
      if (count < datafrac) continue;
      count = 0;
      istringstream linein(line);
      double r,mapsq,mxsq,junk,gamsq;
      double sigmap,siggamsq;
      linein >>r>>mapsq>>sigmap>>mxsq>>junk>>junk>>junk>>junk>>junk>>gamsq>>siggamsq;
      r *= ARCMIN;

#ifdef BMODE
      mapsq += (BMODE)*mxsq;
#endif

#ifdef DILLUT
      mapsq *= DILLUTSQ;
      gamsq *= DILLUTSQ;
      //sigmap *= DILLUTSQ;
      //siggamsq *= DILLUTSQ;
#endif

      if (r > map_minr && r < map_maxr) {
	rmap[k].push_back(r);
	map[k].push_back(mapsq);
	//sigmap[k].push_back(sigmap);
      }
      if (r > var_minr && r < var_maxr) {
	rvar[k].push_back(r);
	var[k].push_back(gamsq);
	//sigvar[k].push_back(sigvar);
      }
    }
    dbg<<"sizes = "<<rmap[k].size()<<", "<<rvar[k].size()<<endl;
    if (k>0) {
      Assert(rmap[k].size() == rmap[0].size());
      Assert(rvar[k].size() == rvar[0].size());
    }
  }
  size_t nmap = rmap[0].size();
  size_t nvar = rvar[0].size();
  size_t ntot = nmap + nvar;
  dbg<<"nmap = "<<nmap<<", nvar = "<<nvar<<endl;

  r = new Vector<double>(ntot,0.);
  mean = new Vector<double>(ntot,0.);
  cov = new Matrix<double>(ntot,ntot,0.);

  for(size_t k=0;k<nfields;k++) {
    for(size_t i=0;i<nmap;i++) {
      (*r)[i] += rmap[k][i];
      (*mean)[i] += map[k][i];
      for(size_t j=0;j<=i;j++) {
	(*cov)(i,j) += map[k][i] * map[k][j];
      }
    }
    for(size_t i=0;i<nvar;i++) {
      (*r)[nmap+i] += rvar[k][i];
      (*mean)[nmap+i] += var[k][i];
      for(size_t j=0;j<=i;j++) {
	(*cov)(nmap+i,nmap+j) += var[k][i] * var[k][j];
      }
      for(size_t j=0;j<nmap;j++) {
	(*cov)(nmap+i,j) += var[k][i]*map[k][j];
      }
    }
  }
  for(size_t i=1;i<ntot;i++) cov->diag(i) = cov->diag(-i);

  *r /= double(nfields);
  *mean /= double(nfields);
  *cov -= double(nfields)*(*mean ^ *mean);
  *cov /= double(nfields*(nfields-1));
  for(size_t i=0;i<ntot;i++) Assert((*cov)(i,i) >= 0.);

#ifdef BMODE
  *cov /= sqrt(BMODE*BMODE+1);
#endif

  nmapout = nmap;

}

#ifdef NonFlat
typedef map<UParamsNonFlat,double> umap;
#else
typedef map<UParams,double> umap;
#endif


void ReadCMB(umap& lnlike)
{
  dbg<<"Start ReadCMB\n";

  dbg<<"Using "<<cmbfile<<" for priors\n";
  ifstream fin(cmbfile);
  if (!fin) myerror("opening cmb file:  ",cmbfile);

  double ombh2,omch2,theta,n,Z,w=-1.,As,h,like;

  // I ignore theta and Z. 
  // ll is apparently not acually the log likelihood.
  // Licia says to ignore it and take the likelihood
  // to be proportional to the repetition of the lines.

#ifdef WMAP3
  double junk;
  //                   -ln(lik)                  theta    tau     om_k
  while (fin >> like >> junk >> ombh2 >> omch2 >> junk >> junk >> junk >>
      //  f_nu              n_t     dn_d/dlnk      r      A_sz    om_lam
      junk >> w >> n >> junk >> junk >> As >> junk >> junk >> junk >>
      //  age     om_m    sig_8   z_rei   ZERO       beta_sdss b_sdss  z_bg
      junk >> junk >> junk >> junk >> junk >> h >> junk >> junk >> junk)
#else
  while (fin >> ombh2 >> omch2 >> theta >> n >> Z >> 
#ifndef LAMBDA
      w >> 
#endif
      As >> h >> like) 
#endif
  {
#ifdef WMAP3
    h /= 100.;
    As = exp(As);
    As /= 29.5;
#endif

#ifdef Wa
    double wa = Z;
    // For Wa, I do set "like" value to be 1/sigma_w^2,
    // which is the relative likelihood for that (w,wa) value
#else
    double wa = 0.;
#ifndef WMAP3
    like = 1.;
    // For WMAP1, just use unit likelihood for each line.
#endif
#endif

#ifdef NonFlat 
    double omde = theta;
    UParamsNonFlat up(ombh2,omch2,omde,n,w,wa,As,h);
#else
    UParams up(ombh2,omch2,n,w,wa,As,h);
#endif

    lnlike[up]+=like;  // for now, _not_ log
  }
  dbg<<"cmblnlike size = "<<lnlike.size()<<endl;
  umap::iterator mapit;
  for(mapit = lnlike.begin(); mapit != lnlike.end(); mapit++) {
    mapit->second = log(mapit->second);
    //dbg<<mapit->first.w<<"  "<<mapit->first.As<<"  "<<mapit->second<<endl;
  }
}

double PredInterp(double r, Vector<double> vlogr, Vector<double> vlogy)
{
  size_t n = vlogr.size();
  Assert(vlogy.size() == n);
  Assert(n >= 2);

  double logr1,logr2,logy1,logy2;
  double logr = log(r);
  if (logr < vlogr[0]) {
    dbg<<"Warning: extrapolating "<<r<<" from "<<exp(vlogr[0])<<", "<<exp(vlogr[1])<<endl;
    logy1 = vlogy[0];
    logy2 = vlogy[1];
    logr1 = vlogr[0];
    logr2 = vlogr[1];
  } else if (logr > vlogr[n-1]) {
    dbg<<"Warning: extrapolating "<<r<<" from "<<exp(vlogr[n-2])<<", "<<exp(vlogr[n-1])<<endl;
    logy1 = vlogy[n-2];
    logy2 = vlogy[n-1];
    logr1 = vlogr[n-2];
    logr2 = vlogr[n-1];
  } else {
    size_t i;
    for(i = 1; logr > vlogr[i]; i++) {
      Assert(i < n-1);
    }
    Assert(i < n);
    logy1 = vlogy[i-1];
    logy2 = vlogy[i];
    logr1 = vlogr[i-1];
    logr2 = vlogr[i];
  }

  double logy = ((logr-logr1)*logy2 + (logr2-logr)*logy1)/(logr2-logr1);
  return exp(logy);
}

int main(int argc, char* argv[]) 
{
#ifdef NAME
  string debugname = string("predlike.") + NAME + ".debug";
#else
  string debugname = string("predlike.debug");
#endif
  dbgout = new ofstream(debugname.c_str());
  //dbgout = &std::cout;

#ifdef DILLUT
  dbg<<"DILLUT = "<<DILLUT<<endl;
#endif
#ifdef BMODE
  dbg<<"BMODE = "<<BMODE<<endl;
#endif
  dbg<<"map range = "<<map_minr/ARCMIN<<" - "<<map_maxr/ARCMIN<<endl;
  dbg<<"var range = "<<var_minr/ARCMIN<<" - "<<var_maxr/ARCMIN<<endl;


  umap cmblnlike;
  ReadCMB(cmblnlike);

  Vector<double>* r;
  Vector<double>* data;
  Matrix<double>* Sigma;
  size_t nmap;
  ReadData(r,data,Sigma,nmap);
  size_t ntot = r->size();

  Sigma->divideUsing(tmv::SV);
  Sigma->saveDiv();
  Sigma->setDiv();
  Sigma->svd().thresh(singthresh,dbgout);

  std::ifstream predin(predfile);
  if (!predin) myerror("opening pred file:  ",predfile);
  dbg<<"opened predfile "<<predfile<<endl;
  double pred_map_minr, pred_map_maxr, pred_map_step;
  double pred_var_minr, pred_var_maxr, pred_var_step;
  int npred_map, npred_var;

  predin >> pred_map_minr >> pred_map_maxr >> npred_map >> pred_map_step;
  dbg<<"pred read: "<<pred_map_minr<<", "<<pred_map_maxr<<", "<<npred_map<<", "<<pred_map_step<<endl;;
  predin >> pred_var_minr >> pred_var_maxr >> npred_var >> pred_var_step;
  dbg<<"pred read: "<<pred_var_minr<<", "<<pred_var_maxr<<", "<<npred_var<<", "<<pred_var_step<<endl;;
  if (!predin) myerror("reading predin ",predfile);
  Vector<double> pred_maplogr(npred_map+1);
  Vector<double> pred_logmap(npred_map+1);
  Vector<double> pred_varlogr(npred_var+1);
  Vector<double> pred_logvar(npred_var+1);
  double pred_map_logminr = log(pred_map_minr);
  double pred_var_logminr = log(pred_var_minr);
  double pred_map_logstep = log(pred_map_step);
  double pred_var_logstep = log(pred_var_step);
  xdbg<<"mapminr = "<<pred_map_minr<<" = "<<exp(pred_map_logminr)<<endl;
  xdbg<<"mapstep = "<<pred_map_step<<" = "<<exp(pred_map_logstep)<<endl;
  xdbg<<"varminr = "<<pred_var_minr<<" = "<<exp(pred_var_logminr)<<endl;
  xdbg<<"varstep = "<<pred_var_step<<" = "<<exp(pred_var_logstep)<<endl;

  umap totlnlike;
  umap::iterator mapit;
#ifdef NAME
  string datname = string("lnlike.") + NAME + ".dat";
#else 
  string datname = string("lnlike.dat");
#endif
  ofstream fout(datname.c_str());
#ifdef Wa
  fout <<"#  om_m    om_lam     om_b        h        As       n         w       wa       sig8     datall      ll\n";
#else
  fout <<"#  om_m    om_lam     om_b        h        As       n         w       sig8     datall      ll\n";
#endif

  int count=0;
  for(mapit=cmblnlike.begin();mapit!=cmblnlike.end();mapit++,count++) {
    dbg<<count<<endl;
#ifdef TAKEEVERYN
    if (count % TAKEEVERYN != 0) continue;
#endif
    double logr = pred_map_logminr;
    for(int i=0;i<=npred_map;i++,logr+=pred_map_logstep) {
      double x;
      predin >> x;
      pred_logmap[i] = log(x);
      pred_maplogr[i] = logr;
      //dbg<<"logr = "<<logr<<", r = "<<exp(logr)<<endl;
      dbg<<"map("<<exp(logr)/ARCMIN<<") = "<<x<<endl;
    }
    logr = pred_var_logminr;
    for(int i=0;i<=npred_var;i++,logr+=pred_var_logstep) {
      double x;
      predin >> x;
      pred_logvar[i] = log(x);
      pred_varlogr[i] = logr;
      dbg<<"var("<<exp(logr)/ARCMIN<<") = "<<x<<endl;
    }
    double sig8;
    predin >> sig8;
    dbg<<"sig8 = "<<sig8<<endl;
    if (!predin && !predin.eof()) myerror("reading predin ",predfile);

#ifdef STARTAT
    if (count < STARTAT) continue;
#endif
#ifdef STOPAT
    if (count > STOPAT) break;
#endif

#ifdef NonFlat
    const UParamsNonFlat& up = mapit->first;
#else
    const UParams& up = mapit->first;
#endif

    double cmbll = mapit->second;
    dbg<<count<<"/"<<cmblnlike.size()<<"  "<<up.ombh2<<"  "<<up.omch2<<"  "<<up.n<<"  "<<up.w<<"  "<<up.wa<<"  "<<up.As<<"  "<<up.h<<endl;

    double datall;
    Vector<double> pred(ntot);

    for(size_t i=0;i<ntot;i++) {
      if (i < nmap) {
	pred[i] = PredInterp((*r)[i],pred_maplogr,pred_logmap);
      } else {
	pred[i] = PredInterp((*r)[i],pred_varlogr,pred_logvar);
      }
      //dbg << (*r)[i]/ARCMIN <<"  "<< (*data)[i] <<"  "<< pred[i] <<"  "<< ((*data)[i]-pred[i])/sqrt((*Sigma(i,i))) <<endl;
    }

    Vector<double> diff = *data - pred;
    //dbg <<"diff = "<<diff<<endl;
    //dbg <<"diff/sig = "<<(diff/(*Sigma))<<endl;
    //dbg <<"diff/sig*sig = "<<(diff/(*Sigma) * (*Sigma))<<endl;
    double chisq = diff * (diff/(*Sigma));
    dbg <<"chisq = "<<chisq<<endl;
    //dbg <<"likelihood = "<<exp(-chisq/2.)<<endl;
    datall = -chisq/2.;

    double ll = datall+cmbll;
    totlnlike[up] = ll;
    double omm = up.GetOmM();
    double omlam = up.GetOmDE();
    double omb = up.GetOmB();
    dbg <<omm<<"  "<<omlam<<"  "<<omb<<"  "<<up.h<<"  "<<up.As<<"  "<<up.n<<"  "<<up.w<<"  "<<up.wa<<"  "<<sig8<<"  "<<datall<<"  "<<ll<<endl;
    dbg <<"totlnlike = "<<datall<<" + "<<cmbll<<" = "<<ll<<endl;
#ifdef Wa
    fout <<omm<<"  "<<omlam<<"  "<<omb<<"  "<<up.h<<"  "<<up.As<<"  "<<up.n<<"  "<<up.w<<"  "<<up.wa<<"  "<<sig8<<"  "<<datall<<"  "<<ll<<endl;
#else
    fout <<omm<<"  "<<omlam<<"  "<<omb<<"  "<<up.h<<"  "<<up.As<<"  "<<up.n<<"  "<<up.w<<"  "<<sig8<<"  "<<datall<<"  "<<ll<<endl;
#endif
  }

  return 0;
}
